#
# This script requires Python 2.6 or newer
#

from argParser import *

from cogent import *
import math
from decimal import Decimal
import os

#
# PREFERENCES:
#
clustalw_executable = "~/Applications/clustalw2"
muscle_executable = "muscle"
gap_opening_costs = [1.0,5.0,7.5,9.7,9.8,9.9,10.0,10.1,10.2,10.3,12.5,15.0]
gap_extension_costs = [0.2] #[0.02, 0.2, 2.0]
unaligned_sequence_file = "sequences.fasta"
#phyml_executable = "/Users/victor/Applications/phyml_v2.4.4/exe/phyml_macOSX"
phyml_executable = "phyml"

#
# GLOBALS:
#
ap = ArgParser(sys.argv)
ALIGNMENTPATH = ap.getArg("--alignment")
output_path = ALIGNMENTPATH + ".elided"
alin = LoadSeqs(ALIGNMENTPATH, aligned=False)
taxa_names = alin.getSeqNames()

list_of_fasta_alignments = [] # a list of filenames for each FASTA alignment we create.  All the files in this list will be elided.
list_of_phylip_alignments = []
fasta_filename_to_phylip_filename = {}
phylip_file_gap_cost = {} # key = alignment filename, value = gap cost for this alignment
gap_cost_lnl = {} # key = tree filename, value = log likelihood of this tree
gap_cost_likelihood = {}
gap_cost_mean_likelihood = {}
gap_cost_adjusted_mean_likelihood = {}
gap_cost_pp = {}

# INPUT:
#     input_sequence_file = an unaligned fasta file
# OUTPUT:
#    several aligned FASTA files will be written, each using different gap opening and
#    extension costs.  Each file's name will be prefixed with <output_alignment_filename_seed>
def run_clustal(input_sequence_file, output_alignment_filename_seed, test = False):
    for gap_opening in gap_opening_costs:
        for gap_extension in gap_extension_costs:
            command = clustalw_executable
            command = command + " -INFILE=" + input_sequence_file
            output_filename =  output_alignment_filename_seed + ".go=" + gap_opening.__str__() + ".gx=" + gap_extension.__str__() + ".fasta"
            command = command + " -OUTFILE=" + output_filename
            list_of_fasta_alignments.append(output_filename)
            command = command + " -PWMATRIX=GONNET"
            command = command + " -GAPOPEN=" + gap_opening.__str__()
            command = command + " -GAPEXT=" + gap_extension.__str__()
            command = command + " -TYPE=PROTEIN"
            command = command + " -OUTPUT=FASTA" 
            if test == False:
                os.system(command) 
            print("Writing Phylip version...")
            fout = open(output_alignment_filename_seed + ".go=" + gap_opening.__str__() + ".gx=" + gap_extension.__str__() + ".phy", "w")
            seqs = LoadSeqs(output_filename)
            fout.write( seqs.toPhylip()[0] + "\n")
            fout.close()
            list_of_phylip_alignments.append(output_alignment_filename_seed + ".go=" + gap_opening.__str__() + ".gx=" + gap_extension.__str__() + ".phy")
            fasta_filename_to_phylip_filename[output_filename] = output_alignment_filename_seed + ".go=" + gap_opening.__str__() + ".gx=" + gap_extension.__str__() + ".phy"
            phylip_file_gap_cost[output_alignment_filename_seed + ".go=" + gap_opening.__str__() + ".gx=" + gap_extension.__str__() + ".phy"] = gap_opening
    print("List of Phylip files = ", list_of_phylip_alignments)

# INPUT:
#     none
#    but, assumes that list_of_fasta_alignments contains a valid list of alignments
# OUTPUT:
#    an elided FASTA alignment will be written to <output_alignment_filename>.fasta
#    and a Phylip version will be written to <output_alignment_filename>.phy
def elide_gene_alignments(output_alignment_filename_seed):
    taxa_sequence = {} # key = gene name, value = elided sequences
    for t in taxa_names:
        taxa_sequence[t] = ""

    # elide the alignments into one mega-alignment    
    for f in list_of_fasta_alignments:
        print("Adding sequences from file: " + f)
        aln = LoadSeqs(f)
        x = aln.toPhylip()
        for line in x[0].split("\n"):
            if line.__len__() < aln.__len__():
                continue
            tokens = line.split()
            taxaid = x[1][ tokens[0] ]
            seq = tokens[1]
            #print seq
            if taxa_sequence.has_key(taxaid):
                taxa_sequence[taxaid] = taxa_sequence[taxaid] + seq
    
    # write the elided alignment to disk
    fout = open(output_alignment_filename_seed + ".fasta", "w")
    for t in taxa_names:
        fout.write(">" + t + "\n")
        fout.write( taxa_sequence[t] + "\n")
    fout.close()

    # create a Phylip version for ProtTest
    # here we just convert the elided file (the one we just wrote) into a Phylip file
    print("Writing Phylip version...")
    fout = open(output_alignment_filename_seed + ".phy", "w")
    seqs = LoadSeqs(output_alignment_filename_seed + ".fasta")
    fout.write( seqs.toPhylip()[0] + "\n")
    fout.close()

# writes a file named [gene_name].phyml.script
# returns the path to the script
#
# model = a model name that is recognized by PhyML
# useinvar = set a fraction of the sites to be invariable
# usegamma = use a gamma-distributed rate of mutation
# useemp = emprically estimate the equilibrium frequencies from the data
#   (as of now, useemp is not actually used)
def writePhymlScript(alignment_file, model, useinvar, usegamma, useemp, tree = "BIONJ"):
    ##################################################
    # Copied From the PHYML command-line manual:
    # Type 12 parameters for amino-acids sequences :
    # 1. sequences file 
    # 2. data type 
    # 3. sequence format 
    # 4. nb data sets 
    # 5. nb bootstrapped data sets 
    # 6. substitution model 
    # 7. prop. invariable sites 
    # 8. nb categories 
    # 9. gamma parameter 
    # 10. starting tree 
    # 11. optimise topology 
    # 12. optimise branch lengths and rate parameters
    # Example :
    #./phyml seqs2 1 i 1 0 JTT 0.0 4 1.0 BIONJ n n
    ####################################################
    scriptpath = alignment_file + ".phyml.script"
    f = open(scriptpath, "w")
    f.write(phyml_executable)
    f.write(" " + alignment_file) #1.
    f.write(" 1") #2: amino acids
    f.write(" i") #3: interleaved data format
    f.write(" 1") #4. 
    f.write(" 0") #5.
    f.write(" " + model) #6.
    if useinvar:
        f.write(" e") #7. use the ML proportion of invariant sites
    else:
        f.write(" 0.0") #7.
    if usegamma:
        f.write(" 4") #8. use 4 gamma categories
    else:
        f.write(" 1") #8.
    f.write(" e") #9. use the ML alpha for gamma
    f.write(" " + tree) #10.
    if tree == "BIONJ":
        f.write(" y") #11.
    else:
        f.write(" n") #11.
    f.write(" n") #12.
    f.write("\n")
    f.close()
    return scriptpath 

def build_elided_tree():
    scriptpath = writePhymlScript(output_path + ".phy", "JTT", False, True, True)
    os.system("source " + scriptpath)

#
# Fills the global hash 'gap_cost_lnl'
#
def calculate_tree_likelihood_per_alignment():
    for f in list_of_phylip_alignments:
        print("Calculating tree likelihood for Phylip alignment ", f)
        # Run PhyML
        scriptpath = writePhymlScript(f, "JTT", False, True, True, tree = output_path + ".phy_phyml_tree.txt")
        os.system("source " + scriptpath)
        # Get tree likelihood
        fin = open(f + "_phyml_stat.txt", "r")
        lines = fin.readlines()
        for l in lines:
            print(l)
            if None != re.match(".*Likelihood : loglk", l):
                tokens = l.split()
                lnl = float(tokens[5])
                print("Lnl = ", lnl)
                gap_cost_lnl[ phylip_file_gap_cost[f] ] = lnl
        fin.close()

#
# Normalizes a set of log likelihood values to manageable numbers (which aren't super tiny)
#
def lnls_to_likelihoods():
    for f in list_of_phylip_alignments:
        # Here I use the decimal.Decimal object.
        # I realize it requires a major performance hit.  However, it allows
        # me to directly operate with likelihood values.
        this_lnl = Decimal( gap_cost_lnl[ phylip_file_gap_cost[f] ].__str__() )
        this_likelihood = this_lnl.exp()
        gap_cost_likelihood[ phylip_file_gap_cost[f] ] = this_likelihood # e^lnl

#
# Fills 'gap_cost_mean_likelihood'
#
def average_likelihoods_per_site():
    # in this method, i'm iterating over FASTA filenames, not PHYLIP filenames.
    # therefore i must lookup the associated phylip filename for every fasta filename.
    # consequently, i use some additional inderection in the hash key lookups.
    # sorry for the obfuscation!
    for f in list_of_fasta_alignments:
        # how many sites were used in this alignment?
        count_non_gap_sites = 0
        aln = LoadSeqs(f)
        for site in range(0, aln.__len__()):
            if False == aln[site].columnFreqs()[0].keys().__contains__("-"):
                count_non_gap_sites += 1
        print "gap cost: ", phylip_file_gap_cost[ fasta_filename_to_phylip_filename[f] ], " non-gap sites=", count_non_gap_sites
        s = float(1) / float(count_non_gap_sites)
        s = Decimal( s.__str__() )
        gap_cost_mean_likelihood[ phylip_file_gap_cost[ fasta_filename_to_phylip_filename[f] ] ] = gap_cost_likelihood[phylip_file_gap_cost[ fasta_filename_to_phylip_filename[f] ]]**s
        gap_cost_adjusted_mean_likelihood[ phylip_file_gap_cost[ fasta_filename_to_phylip_filename[f] ] ] = gap_cost_mean_likelihood[phylip_file_gap_cost[ fasta_filename_to_phylip_filename[f] ]] * s  
      

#
# Given a set of log-likelihood values, calculates a posterior probability distribution
#
def lnl_to_pp():
    sum = Decimal(0)
    for f in list_of_phylip_alignments:
        sum += gap_cost_adjusted_mean_likelihood[ phylip_file_gap_cost[f] ]
        
    
    # 4. calculate post probs:
    gap_cost_pp = {}
    for f in list_of_phylip_alignments:
        gap_cost_pp[ phylip_file_gap_cost[f] ] = gap_cost_adjusted_mean_likelihood[ phylip_file_gap_cost[f] ] / sum
        
# where keys are X values and hash values are Y values.
def plot_hash(h, title):
    filename_seed = re.sub(" ", "_", title)
    xpoints = []
    ypoints = []
    for item in h:
        xpoints.append( float(item) )
        ypoints.append( float(h[item]) )
    f = open(filename_seed + ".cran", "w")
    x = "x <- c("
    y = "y <- c("
    for i in range(0, xpoints.__len__() ):
        x += xpoints[i].__str__() + ","
        y += ypoints[i].__str__() + ","
    x = re.sub(",$", "", x)
    x += ")"
    y = re.sub(",$", "", y)
    y += ")"
    f.write( y + "\n")
    f.write( x + "\n")    
    f.write("pdf(\"" + filename_seed + ".pdf" + "\") \n")
    f.write("plot(x, y, type=\"b\", main=\"" + title + "\", pch=20)\n")
    f.write("dev.off()\n")
    f.close()
    print "Dispatching to R (be patient). . ."
    os.system("R --save < " + filename_seed + ".cran > " + filename_seed + ".log")    

################################################################
# MAIN:
#
run_clustal(ALIGNMENTPATH, output_path, test=False) # 'test' will prevent the method from actually invoking clustal
elide_gene_alignments(output_path)
build_elided_tree()
calculate_tree_likelihood_per_alignment()
print "gap_cost_lnl: ", gap_cost_lnl
lnls_to_likelihoods()
print "gap_cost_likelihood: ", gap_cost_likelihood
average_likelihoods_per_site()
print "gap_cost_mean_likelihood: ", gap_cost_mean_likelihood
print "gap_cost_adjusted_mean_likelihood: ", gap_cost_adjusted_mean_likelihood   
lnl_to_pp()
print "gap_cost_pp: ", gap_cost_pp

#
# my next question: how do the absolute number of gaps vs. gap cost look?
#

plot_hash(gap_cost_lnl, "gap cost v lnL")
plot_hash(gap_cost_likelihood, "gap cost v likelihood")
plot_hash(gap_cost_mean_likelihood, "gap cost v mean likelihood of a site")
plot_hash(gap_cost_adjusted_mean_likelihood, "gap cost v adjusted mean likelihood of a site")
plot_hash(gap_cost_pp, "gap cost v PP")


 
    